package LeetCode;

import java.util.*;

public class LC_355_DesignTwitter {

    public static void main(String[] args) {

    }

    class Wrapper {
        ArrayList<Integer> list;
        int index;

        Wrapper(ArrayList<Integer> list, int index) {
            this.list = list;
            this.index = index;
        }
    }

    public class Twitter {
        HashMap<Integer, HashSet<Integer>> userMap;
        HashMap<Integer, ArrayList<Integer>> tweetMap;
        HashMap<Integer, Integer> orderMap;
        int order;

        public Twitter() {
            userMap = new HashMap<>();
            tweetMap = new HashMap<>();
            orderMap = new HashMap<>();
        }

        public void postTweet(int userId, int tweetId) {
            ArrayList<Integer> list = tweetMap.computeIfAbsent(userId, k -> new ArrayList<>());
            list.add(tweetId);
            orderMap.put(tweetId, order++);
            follow(userId, userId);
        }

        public List<Integer> getNewsFeed(int userId) {
            HashSet<Integer> set = userMap.get(userId);
            if (set == null) return new ArrayList<>();
            ArrayList<ArrayList<Integer>> lists = new ArrayList<>();
            for (int uid : set) {
                if (tweetMap.get(uid) != null && tweetMap.get(uid).size() > 0)
                    lists.add(tweetMap.get(uid));
            }
            ArrayList<Integer> result = new ArrayList<>();
            PriorityQueue<Wrapper> queue = new PriorityQueue<>((a, b) -> orderMap.get(b.list.get(b.index)) - orderMap.get(a.list.get(a.index)));
            for (ArrayList<Integer> list : lists) queue.offer(new Wrapper(list, list.size() - 1));
            while (!queue.isEmpty() && result.size() < 10) {
                Wrapper top = queue.poll();
                result.add(top.list.get(top.index));
                top.index--;
                if (top.index >= 0) queue.offer(top);
            }
            return result;
        }

        void follow(int followerId, int followeeId) {
            HashSet<Integer> set = userMap.computeIfAbsent(followerId, k -> new HashSet<>());
            set.add(followeeId);
        }

        public void unfollow(int followerId, int followeeId) {
            if (followerId == followeeId) return;
            HashSet<Integer> set = userMap.get(followerId);
            if (set == null) return;
            set.remove(followeeId);
        }
    }

}